package scfg.utils;

import java.util.*;
import java.util.Map.Entry;
import java.awt.font.ImageGraphicAttribute;
import java.io.*;

import scfg.output.Display;

public class alnFastaParser {
	
	
	public static boolean IGNORE_ACC_EXTENSION;
	public static boolean DEBUG_ACCESSION_NUM;
	
	private static Display output;
	private static Display html;
	private static Display ambiguous;
	private static Display duplicateAccNum;
	private static Display wrongFormat;
	private static int tempNothingFound;
	private static int tempWrongSeq;
	private static int tempTotalFiles;
	private static int tempExactMatches;
	private static int tempHtmlCount;
	private static int tempAmbiguousCount;
	private static Display tempRM;
	private static Map<String, List<String>> uniqueSequences;
	
	static {
		output = new Display("alnFastaParser", false);
		ambiguous = new Display("alnFastaParser_ambiguous", false);
		html = new Display("html_files", false);
		duplicateAccNum = new Display("alnFastaParser_AccNum", false);
		wrongFormat = new Display("ctFiles_wrongFormat", false);
		tempNothingFound = 0;
		tempWrongSeq = 0;
		IGNORE_ACC_EXTENSION = false;
		DEBUG_ACCESSION_NUM = false;
		uniqueSequences = new HashMap<String, List<String>>();
		tempRM = new Display("temp_remove", false);
	}
	
	private static void resetOutputs(String prepend) {
		output = new Display(prepend + "alnFastaParser", false);
//		ambiguous = new Display(prepend + "alnFastaParser_ambiguous", false);
		duplicateAccNum = new Display(prepend + "alnFastaParser_AccNum", false);
	}

	public static HashMap<String, List<String>> parseFastaFile(String filename) {
		return parseFastaFile(filename, new HashMap<String, List<String>>());
	}
	
	public static HashMap<String, List<String>> parseFastaFile(String filename, HashMap<String, List<String>> alignments) {
		if (filename == null)
			return new HashMap<String, List<String>>();
		File fastaFile = new File(filename);
		if (!fastaFile.exists())
			return alignments;

		try {
			Scanner scan = new Scanner(fastaFile);
			String line = scan.nextLine();
			int count = 0;
			while (scan.hasNextLine()) {
				if ((count++)%20==0)
					System.out.println(alignments.size());
				if (line.startsWith(">"))
					line = handleAlignment(line, scan, alignments);
				else {
					// TODO: There is a problem
					line = scan.nextLine();
				}
			}
			scan.close();
		} catch (Exception e) {
			// TODO: report exception
		}
		return alignments;
	}

	private static String handleAlignment(String line, Scanner scan,
			HashMap<String, List<String>> alignments) {
		String[] arr = line.split("::");
		if (arr.length == 3 && arr[2].length() > 1) {
			String AccessionNum;
			if (IGNORE_ACC_EXTENSION) {
				arr = arr[2].split("\\.");
				AccessionNum = arr[0].trim();
			} else
				AccessionNum = arr[2].trim();
			if (AccessionNum.equals("DIVIDER")){
//				output.out("DIVIDER (no accession num)");
				return scan.nextLine();
			}
			StringBuilder align = new StringBuilder();
			while (scan.hasNextLine()) {
				line = scan.nextLine();
				if (line.startsWith(">"))
					break;
				else
					align.append(line);
			}
			String alignStr = align.toString().toUpperCase();
			if (alignStr.toUpperCase().matches("[ACGU-]+") && !AccessionNum.equals("DIVIDER")){
//				if(alignments.get(AccessionNum) != null && !alignments.get(AccessionNum).equals(alignStr)){
////					output.out("There is a big problem... Same accession number, different sequences.");
//					duplicateAccNum.eout("Duplicate Accession Number w/ non-matching sequence :: " + AccessionNum + " :: " + alignStr);
////					output.out("\t[ X ] " + AccessionNum + " :: " + alignments.get(AccessionNum));
//				}
//				alignments.put(AccessionNum, alignStr);
				List<String> temp = alignments.get(AccessionNum);
				if (temp == null){
					temp = new LinkedList<String>();
					alignments.put(AccessionNum, temp);
				}
				if (!temp.contains(alignStr))
					temp.add(alignStr);
			} else {
//				ambiguous.eout("Ambiguous :: " + AccessionNum + " :: " + alignStr);
			}
			return line;
		}
		// TODO: There is a problem
		return scan.next();
	}
	
	public static List<List<String>> findApprovedCTFiles(String dirName, HashMap<String, List<String>> alignments) {
		if (dirName == null)
			return null;
		return findApprovedCTFiles(new File(dirName), alignments);
	}
	
	public static List<List<String>> findApprovedCTFiles(File base, HashMap<String, List<String>> alignments){
		if (base == null || alignments == null || !base.exists() || !base.isDirectory())
			return null;
		
		List<List<String>> filenames = new ArrayList<List<String>>();
		for(int i=0;i<3;i++)
			filenames.add( new LinkedList<String>());
		
		List<File> dirs = new ArrayList<File>();
		File[] files = base.listFiles();
		int count = 0;
		for (File f : files){
			if(count++ % 500 == 0)
				System.out.println((count-1) + "/" + files.length);
			if (f.isDirectory())
				dirs.add(f);
			else {
				Boolean test = testApprovedCTFile(f, alignments);
				if (test == null)
					filenames.get(2).add(f.getAbsolutePath());
				else if(test)
					filenames.get(0).add(f.getAbsolutePath());
				else
					filenames.get(1).add(f.getAbsolutePath());
			}
		}
		
		for (File d : dirs)
			findApprovedCTFiles(d, alignments);
		
		return filenames;
	}
	
	private static Boolean testApprovedCTFile(File f, HashMap<String, List<String>> alignments){
		if (f == null || alignments == null || !f.exists() || !f.getName().contains(".nopct")){
			if (f.getName().contains(".rna")){
				
			}
			return null;
		}
		try {
			Scanner scan = new Scanner(f);
			String line = scan.nextLine();
			if (line.contains("html")){
				html.out(f.getAbsolutePath());
				tempHtmlCount++;
				return null;
			}
			if (!line.matches("Filename:.*")){
				wrongFormat.out(f.getAbsolutePath());
				return null;
			}
			tempRM.outC("rm ").outC(f.getName()).out(".rna");
			tempTotalFiles++;
			scan.nextLine();
			String[] AccessionNumbers = scan.nextLine().replaceAll("Accession Numbers?:\\s*", "").replaceAll(",", "").split("\\s+");
			if (DEBUG_ACCESSION_NUM && AccessionNumbers.length > 1)
				System.out.println(AccessionNumbers.length + " :: " + f.getName() + " :: " + Arrays.toString(AccessionNumbers));
			if (IGNORE_ACC_EXTENSION) {
				for(int i=0;i<AccessionNumbers.length;i++) {
					String[] arr = AccessionNumbers[i].split("\\.");
					AccessionNumbers[i] = arr[0].trim();
				}
			}
			StringBuilder sb = new StringBuilder();
			while(scan.hasNextLine()){
				String[] arr = scan.nextLine().replaceAll("\\s+", " ").split("\\s");
				if (arr.length == 7 && arr[2].matches("[A-Z]")) {
					sb.append(arr[2].toUpperCase());
				}
			}
			scan.close();
			String ctSeq = sb.toString();
			if (!ctSeq.toUpperCase().matches("[ACGU]+")) {
				ambiguous.out(f.getAbsolutePath());
				tempAmbiguousCount++;
				return null;
			}
			boolean nothingFound = true;
			for (String AccessionNumber : AccessionNumbers) {
				List<String> seqList = alignments.get(AccessionNumber);
				if (seqList != null) {
					nothingFound = false;
					for (String seq : seqList) {
						seq = seq.replaceAll("-", "").toUpperCase();
						if(seq.equals(ctSeq)){
							tempExactMatches++;
							// TODO: unique sequences
							if (uniqueSequences.get(ctSeq) == null)
								uniqueSequences.put(ctSeq, new LinkedList<String>());
							uniqueSequences.get(ctSeq).add(f.getName());
							return true;
						}
					}
				}
			}
			if (nothingFound) {
				tempNothingFound++;
				if (DEBUG_ACCESSION_NUM)
					System.out.println("NF: " + AccessionNumbers.length + " :: " + Arrays.toString(AccessionNumbers));
			} else {
				tempWrongSeq++;
				if (DEBUG_ACCESSION_NUM)
					System.out.println("WS: " + AccessionNumbers.length + " :: " + Arrays.toString(AccessionNumbers) + " :: " + ctSeq);
			}
			return null;
		} catch (Exception e) {
			// TODO: 
		}
		return null;
	}
	
	public static boolean sortCTWithFasta(String ctDirName, String fastaDirName, boolean verbose){
		return sortCTWithFasta(ctDirName, fastaDirName, verbose, "");
	}
	
	private static boolean sortCTWithFasta(String ctDirName, String fastaDirName, boolean verbose, String prepend){
		return sortCTWithFasta(ctDirName, fastaDirName, verbose, prepend, null);
	}
	
	private static boolean sortCTWithFasta(String ctDirName, String fastaDirName, boolean verbose, String prepend, String alignmentsFilename){
		if (ctDirName == null || fastaDirName == null)
			return false;
		File ctDir = new File(ctDirName);
		File fastaDir = new File(fastaDirName);
		if (!ctDir.exists() || !fastaDir.exists())
			return false;
		HashMap<String, List<String>> alignments = getAlignments(alignmentsFilename, fastaDir.listFiles());
		List<List<String>> filenames = findApprovedCTFiles(ctDirName, alignments);
		output = new Display(prepend + "ct_approved.txt", null, verbose);
//		output.out("Approved:\n---------");
		for (String s : filenames.get(0)){
			String[] arr = s.split("/+");
			output.outC("cp ").outC(s).outC(" acc/").out(arr[arr.length - 1]);
		}
		output.finalizeDisplay();
		output = new Display(prepend + "ct_noMatch.txt", null, verbose);
		output.out("No Match:\n---------");
		for (String s : filenames.get(1))
			output.out(s);
		output.finalizeDisplay();
		output = new Display(prepend + "ct_errorMatch.txt", null, verbose);
		output.out("Error Match:\n------------");
		for (String s : filenames.get(2))
			output.out(s);
		output.finalizeDisplay();
		return true;
	}
	
	private static HashMap<String, List<String>> getAlignments(String alignmentFilename, File[] files) {
		HashMap<String, List<String>> alignments = new HashMap<String, List<String>>();
		if (/* alignmentFilename == null && */ files != null) {
			for(File f : files) {
				if (f.isDirectory()) {
//					dirs.add(f);
				} else 
					parseFastaFile(f.getAbsolutePath(), alignments);
			}
			for (Entry<String, List<String>> me : alignments.entrySet())
				for (String seq : me.getValue())
					output.out(me.getKey() + "::" + seq);
			output.finalizeDisplay();
		} else if (alignmentFilename != null) {
			try {
				Scanner scan = new Scanner(new File(alignmentFilename));
				while(scan.hasNextLine()) {
					String[] arr = scan.nextLine().split("::");
					arr[0] = arr[0].trim();
					if (arr.length == 2) {
						if (alignments.get(arr[0]) == null)
							alignments.put(arr[0], new LinkedList<String>());
						if (!alignments.get(arr[0]).contains(arr[1]))
							alignments.get(arr[0].trim()).add(arr[1].trim());
					}
				}
				scan.close();
			} catch (Exception e) {
				// TODO: 
			}
		}
		return alignments;
	}

	public static void main(String[] args) {	
		long start = System.currentTimeMillis();
		
		ambiguous.out(Display.underline("Ambiguous:", ""));
		
		resetOutputs("5S_");
		String ctDir = "/home/david/Documents/REU/RNA/ct_files/5S";
		String fastaDir = "/home/david/Documents/REU/RNA/fasta_alignments/5S";
		String alignmentFile = "/home/david/Documents/REU/RNA/fasta_alignments/5S/logs/5S_alnFastaParser.log";
//		sortCTWithFasta(ctDir, fastaDir, false, "5S_", alignmentFile);
		
//		output.finalizeDisplay();
//		ambiguous.finalizeDisplay();
//		duplicateAccNum.finalizeDisplay();

		tempDisp();
		
		System.gc();
		
//		resetOutputs("16S_");
//		ctDir = "/home/david/Documents/REU/RNA/ct_files/16S";
//		fastaDir = "/home/david/Documents/REU/RNA/fasta_alignments/16S";
//		alignmentFile = "/home/david/Documents/REU/RNA/fasta_alignments/16S/logs/16S_alnFastaParser.log";
//		sortCTWithFasta(ctDir, fastaDir, false, "16S_", alignmentFile);
		
		
		output.finalizeDisplay();
//		ambiguous.finalizeDisplay();
		duplicateAccNum.finalizeDisplay();

		tempDisp();
		
		System.gc();
		
		resetOutputs("tRNA_");
		ctDir = "/home/david/Documents/RNA/ct_files/tRNA";
		fastaDir = "/home/david/Documents/RNA/fasta_alignments/tRNA";
		alignmentFile = "/home/david/Documents/RNA/fasta_alignments/tRNA/logs/tRNA_alnFastaParser.log";
		sortCTWithFasta(ctDir, fastaDir, false, "16S_", alignmentFile);
		
		
		output.finalizeDisplay();
//		ambiguous.finalizeDisplay();
		duplicateAccNum.finalizeDisplay();

		tempDisp();
		
		System.gc();
		
//		resetOutputs("23S_");
//		ctDir = "/home/david/Documents/REU/RNA/ct_files/23S";
//		fastaDir = "/home/david/Documents/REU/RNA/fasta_alignments/23S";
//		alignmentFile = "/home/david/Documents/REU/RNA/fasta_alignments/23S/logs/23S_alnFastaParser.log";
//		sortCTWithFasta(ctDir, fastaDir, false, "23S_", alignmentFile);
//		
//		output.finalizeDisplay();
		ambiguous.finalizeDisplay();
		duplicateAccNum.finalizeDisplay();
		html.finalizeDisplay();
		tempRM.finalizeDisplay();
		wrongFormat.finalizeDisplay();

		tempDisp();
		
//		System.gc();
		
//		resetOutputs("tRNA_");
//		ctDir = "/home/david/Documents/REU/RNA_big/ct_files/tRNA";
//		fastaDir = "/home/david/Documents/REU/RNA_big/fasta_alignments/tRNA";
//		alignmentFile = "/home/david/Documents/REU/RNA_big/fasta_alignments/tRNA/logs/tRNA_alnFastaParser.log";
//		sortCTWithFasta(ctDir, fastaDir, false, "tRNA_", alignmentFile);
		
//		output.finalizeDisplay();
//		ambiguous.finalizeDisplay();
//		duplicateAccNum.finalizeDisplay();

//		tempDisp();
		
		System.out.println("\n\nfinished: " + ((System.currentTimeMillis() - start) / 1000) + "s");
	}
	
	private static void tempDisp(){
		System.out.println("\nTotal Files: " + tempTotalFiles);
		System.out.println("HTML Count: " + tempHtmlCount);
		System.out.println("Exact Matches: " + tempExactMatches);
		System.out.println("Unique Exact Matches: " + uniqueSequences.size());
		System.out.println("Nothing Found: " + tempNothingFound);
		System.out.println("Wrong Seq: " + tempWrongSeq);
		System.out.println("Ambiguous Count: " + tempAmbiguousCount + "\n\n");
		tempTotalFiles = 0;
		tempExactMatches = 0;
		tempHtmlCount = 0;
		tempNothingFound = 0;
		tempWrongSeq = 0;
		tempAmbiguousCount = 0;
		uniqueSequences.clear();
	}

}
